# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# PyTorch (https://pytorch.org/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
WithEnergyCell
"""

from typing import Union, List, Tuple
import torch
from torch import Tensor
from torch import nn
from torch.nn import Parameter

from ...function import Units, get_arguments
from ...partition import NeighbourList
from ...system import Molecule
from ...potential import PotentialCell
from ...potential.bias import Bias
from ...sampling.wrapper import EnergyWrapper


class WithEnergyCell(nn.Module):
    r"""
    Module that wraps the simulation system with the potential energy function.
    This Module calculates the value of the potential energy of the system at the current coordinates and returns it.

    Args:
        system(Molecule):               Simulation system.
        potential(PotentialCell):       Potential energy function cell.
        bias(Union[Bias, List[Bias]]):  Bias potential function cell. Default: None
        cutoff(float):                  Cut-off distance for neighbour list. If None is given, it will be assigned
                                        as the cutoff value of the of potential energy.
                                        Defulat: None
        neighbour_list(NeighbourList):  Neighbour list. Default: None
        wrapper(EnergyWrapper):         Network to wrap and process potential and bias.
                                        Default: None

    Inputs:
        - **\*inputs** (Tuple(Tensor)) - Tuple of input tensors of 'WithEnergyCell'.

    Outputs:
        energy, Tensor of shape `(B, 1)`. Data type is float. Total potential energy.

    Supported Platforms:
        ``CPU`` ``GPU``

    Symbols:
        B:  Batchsize, i.e. number of walkers of the simulation.
        A:  Number of the atoms in the simulation system.
        N:  Number of the maximum neighbouring atoms.
        U:  Number of potential energy terms.
        V:  Number of bias potential terms.
    """

    def __init__(self,
                 system: Molecule,
                 potential: PotentialCell,
                 bias: Union[Bias, List[Bias]] = None,
                 cutoff: float = None,
                 neighbour_list: NeighbourList = None,
                 wrapper: EnergyWrapper = None,
                 centroid: bool = False,
                 **kwargs
                 ):

        super().__init__()
        self._kwargs = get_arguments(locals(), kwargs)

        self.system = system
        self.potential_function = potential

        self.units = Units(self.system.length_unit,
                           self.potential_function.energy_unit)
        self.system.units.set_energy_unit(self.energy_unit)

        self.bias_function: List[Bias] = None
        self._num_biases = 0
        self._bias_names = []
        if bias is not None:
            if isinstance(bias, list):
                self._num_biases = len(bias)
                self.bias_function = nn.ModuleList(bias)
            elif isinstance(bias, nn.Module):
                self._num_biases = 1
                self.bias_function = nn.ModuleList([bias])
            else:
                raise TypeError(
                    f'The "bias" must be Module or list but got: {type(bias)}')

            for i in range(self._num_biases):
                self._bias_names.append(self.bias_function[i].name)

        self.num_walker = self.system.num_walker
        self.num_atoms = self.system.num_atoms

        self.energy_wrapper = wrapper
        if wrapper is None:
            self.energy_wrapper = EnergyWrapper(
                length_unit=self.length_unit, energy_unit=self.energy_unit)

        self.wrapper_pace = self.energy_wrapper.update_pace

        self.exclude_index = self.potential_function.exclude_index
        self.neighbour_list = neighbour_list
        if neighbour_list is None:
            if cutoff is None and self.potential_function.cutoff is not None:
                cutoff = self.potential_function.cutoff
            self.neighbour_list = NeighbourList(
                system, cutoff, exclude_index=self.exclude_index)
        else:
            self.neighbour_list.set_exclude_index(self.exclude_index)

        self.neighbour_index = self.neighbour_list.neighbours
        self.neighbour_mask = self.neighbour_list.neighbour_mask
        self.num_neighbours = self.neighbour_list.num_neighbours

        if self.neighbour_list.cutoff is not None:
            if self.potential_function.cutoff is None:
                self.potential_function.set_cutoff(
                    self.neighbour_list.cutoff, self.length_unit)
            else:
                pot_cutoff = self.units.length(
                    self.potential_function.cutoff, self.potential_function.length_unit)
                nl_cutoff = self.neighbour_list.cutoff
                if self.potential_function.cutoff > self.neighbour_list.cutoff:
                    raise ValueError(f'The cutoff of the potential function ({pot_cutoff} {self.length_unit}) '
                                     f'cannot be greater than '
                                     f'the cutoff of the neighbour list ({nl_cutoff} {self.length_unit}).')

        self.coordinate = self.system.coordinate
        self.pbc_box = self.system.pbc_box
        self.atom_mass = self.system.atom_mass

        self.pbc_box = self.system.pbc_box

        self.potential_function.set_pbc(self.pbc_box is not None)

        for p in self.potential_function.parameters():
            p.requires_grad = False

        self.potential_function_units = self.potential_function.units

        self.length_unit_scale = torch.tensor(self.units.convert_length_to(
            self.potential_function.length_unit), dtype=torch.float32)

        self._energies = torch.zeros((self.num_walker, self.num_energies),dtype=torch.float32)
        
        self._energy = torch.zeros((self.num_walker, 1),dtype=torch.float32)

        if self.num_biases > 0:
            self._biases = torch.zeros((self.num_walker, self.num_biases),dtype=torch.float32)
            self._bias = torch.zeros((self.num_walker, 1),dtype=torch.float32)
        else:
            self._biases = None
            self._bias = None

        self.centroid = centroid

    @property
    def cutoff(self) -> Tensor:
        r"""
        Cutoff distance for neighbour list.

        Return:
            Tensor, cutoff distance.
        """
        return self.neighbour_list.cutoff

    @property
    def neighbour_list_pace(self) -> int:
        r"""
        Update step for neighbour list.

        Return:
            int, update steps.
        """
        return self.neighbour_list.pace

    @property
    def length_unit(self) -> str:
        r"""
        Length unit.

        Return:
            str, length unit.
        """
        return self.units.length_unit

    @property
    def energy_unit(self) -> str:
        r"""
        Energy unit.

        Return:
            str, energy unit.
        """
        return self.units.energy_unit

    @property
    def num_energies(self) -> int:
        r"""
        Number of energy terms :math:`U`.

        Return:
            int, number of energy terms.
        """
        return self.potential_function.num_energies

    @property
    def num_biases(self) -> int:
        r"""
        Number of bias potential energies :math:`V`.

        Return:
            int, number of bias potential energies.
        """
        return self._num_biases + self.energy_wrapper.num_biases

    @property
    def energy_names(self) -> list:
        r"""
        Names of energy terms.

        Return:
            list[str], names of energy terms.
        """
        return self.potential_function.energy_names

    @property
    def bias_names(self) -> list:
        r"""
        Names of bias potential terms.

        Return:
            list[str], names of bias potential terms.
        """
        bias_names = self._bias_names
        if self.energy_wrapper.num_biases > 0:
            bias_names.extend(self.energy_wrapper.bias_names)
        return bias_names

    @property
    def energies(self) -> Tensor:
        r"""
        Tensor of energy terms in shape (B, U).

        Return:
            Tensor, energy tensor.
        """
        return self.get_energies()

    @property
    def biases(self) -> Tensor:
        r"""
        Tensor of bias potential energies in shape (B, V).

        Return:
            Tensor, bias tensor.
        """
        return self.get_biases()

    @property
    def bias(self) -> Tensor:
        r"""
        Total bias potential in shape (B, 1).

        Return:
            Tensor, bias tensor.
        """
        return self.get_bias()

    def get_energy(self) -> Tensor:
        r"""
        get energy tensor.
        Return:
            Tensor, energy tensor.
        """
        return self._energy
    
    def get_energies(self) -> Tensor:
        r"""
        get energy tensor.

        Return:
            Tensor, energy tensor.
        """
        return self._energies

    def get_biases(self) -> Tensor:
        r"""
        get biases tensor.

        Return:
            Tensor, biases tensor.
        """
        return self._biases

    def get_bias(self) -> Tensor:
        r"""
        get bias tensor.

        Return:
            Tensor, bias tensor.
        """
        return self._bias

    def bias_pace(self, index: int = 0) -> int:
        r"""
        get update step for bias potential term.

        Args:
            index(int):  bias ID. Default: 0.

        Return:
            int, number of update steps.
        """
        return self.bias_function[index].update_pace if self.bias_function is not None else None

    def set_pbc_grad(self, grad_box: bool):
        r"""
        Set whether to calculate gradient for PBC box.

        Args:
            grad_box(bool):  Whether to calculate the gradient for pbc_box.
        """
        self.system.set_pbc_grad(grad_box)

    def update_neighbour_list(self) -> Tuple[Tensor, Tensor]:
        r"""
        Update neighbour list.

        Return:
            tuple[Tensor, Tensor], neighbour_index and neighbour_mask.
        """
        return self.neighbour_list.update(self.coordinate, self.pbc_box)

    def update_bias(self, steps: int):
        r"""
        Update bias potential terms.

        Args:
            step(int):  Current step number.
        """
        if self.bias_function is not None:
            for i in range(self._num_biases):
                if self.bias_pace(i) > 0 and steps % self.bias_pace(i) == 0:
                    self.bias_function[i].update(self.coordinate, self.pbc_box)

    def update_wrapper(self, steps: int):
        r"""
        Update energy wrapper.

        Args:
            step(int):  Current step number.
        """
        if self.wrapper_pace > 0 and steps % self.wrapper_pace == 0:
            self.energy_wrapper.update()
    def get_neighbour_list(self) -> Tuple[Tensor, Tensor]:
        r"""
        Get the neighbour list.

        Return:
            tuple[Tensor, Tensor], neighbour_index and neighbour_mask.
            
            - neighbour_index:  Tensor of shape (B, A, N), data type is int.
            - neighbour_mask:   Tensor of shape (B, A, N), data type is bool.
        """
        return self.neighbour_list.get_neighbour_list()

    def calc_energies(self) -> Tensor:
        r"""
        Calculate energy terms.

        Return:
            Tensor, energy terms in shape (B, U).
        """

        if self.centroid:
            coordinate = self.coordinate.mean(dim=0, keepdim=True)
            pbc_box = self.pbc_box
            if pbc_box is not None:
                pbc_box = pbc_box.mean(dim=0, keepdim=True)
        else:
            coordinate = self.coordinate
            pbc_box = self.pbc_box

        neigh_idx, neigh_vec, neigh_dis, neigh_mask = self.neighbour_list(coordinate, pbc_box)

        coordinate = coordinate * self.length_unit_scale
        if pbc_box is not None:
            pbc_box = pbc_box * self.length_unit_scale

        if neigh_idx is not None:
            neigh_vec = neigh_vec * self.length_unit_scale
            neigh_dis = neigh_dis * self.length_unit_scale

        energies = self.potential_function(
            coordinate=coordinate,
            neighbour_index=neigh_idx,
            neighbour_mask=neigh_mask,
            neighbour_vector=neigh_vec,
            neighbour_distance=neigh_dis,
            pbc_box=pbc_box
        )

        return energies

    def calc_biases(self) -> Tensor:
        r"""
        Calculate bias potential terms.

        Return:
            Tensor, bias potential energy in shape (B, 1).
        """
        if self.bias_function is None:
            return None

        neigh_idx, neigh_vec, neigh_dis, neigh_mask = self.neighbour_list(
            self.coordinate, self.pbc_box)

        coordinate = self.coordinate * self.length_unit_scale
        pbc_box = self.pbc_box
        if pbc_box is not None:
            pbc_box = pbc_box * self.length_unit_scale
        neigh_vec = neigh_vec * self.length_unit_scale
        neigh_dis = neigh_dis * self.length_unit_scale

        biases = []
        for i in range(self._num_biases):
            bias_ = self.bias_function[i](
                coordinate=coordinate,
                neighbour_index=neigh_idx,
                neighbour_mask=neigh_mask,
                neighbour_vector=neigh_vec,
                neighbour_distance=neigh_dis,
                pbc_box=pbc_box
            )
            biases.append(bias_)

        return torch.concat(biases, dim=-1)

    def forward(self, *inputs) -> Tensor:
        r"""
        Calculate the total potential energy of the simulation system.

        Args:
            *inputs:    Inputs of this cell. Typically is empty.

        Return:
            energy(Tensor):  Tensor of shape (B, 1). Data type is float.
                            Total potential energy.

        Symbols:
            B:  Batchsize, i.e. number of walkers of the simulation.
        """

        if self.centroid:
            coordinate = self.system.coordinate.mean(dim=0, keepdim=True)
            pbc_box = self.system.pbc_box
            if pbc_box is not None:
                pbc_box = pbc_box.mean(dim=0, keepdim=True)
        else:
            coordinate = self.system.coordinate
            pbc_box = self.system.pbc_box

        neigh_idx, neigh_mask, neigh_vec, neigh_dis = self.neighbour_list(coordinate, pbc_box)

        coordinate = coordinate * self.length_unit_scale
        if pbc_box is not None:
            pbc_box = pbc_box * self.length_unit_scale

        if neigh_idx is not None:
            neigh_vec = neigh_vec * self.length_unit_scale
            neigh_dis = neigh_dis * self.length_unit_scale

        energies = self.potential_function(
            coordinate=coordinate,
            neighbour_index=neigh_idx,
            neighbour_mask=neigh_mask,
            neighbour_vector=neigh_vec,
            neighbour_distance=neigh_dis,
            pbc_box=pbc_box
        )
        energies = energies.expand(self.num_walker, self.num_energies)
        self._energies.data.copy_(energies)

        biases = None
        if self.bias_function is not None:
            biases = []
            for i in range(self._num_biases):
                bias_ = self.bias_function[i](
                    coordinate=coordinate,
                    neighbour_index=neigh_idx,
                    neighbour_mask=neigh_mask,
                    neighbour_vector=neigh_vec,
                    neighbour_distance=neigh_dis,
                    pbc_box=pbc_box
                )
                biases.append(bias_)
            biases = torch.concat(biases, dim=-1)
            biases = biases.expand(self.num_walker, self.num_biases)

        energy, bias = self.energy_wrapper(energies, biases)
        self._energy.data.copy_(energy)

        if self.energy_wrapper.num_biases > 0:
            if biases is None:
                biases = self.energy_wrapper.get_biases()
            else:
                biases = [biases, self.energy_wrapper.get_biases()]
                biases = torch.concat(biases, axis=-1)
        
        if self.num_biases > 0:
            self._biases.data.copy_(biases)
            self._bias.data.copy_(bias)
        
        return energy
        